import collections
import os

import pandas as pd
import tushare as ts
from loguru import logger

from models.global_config import glm
from models.stock_model import StockNumber, DayInfo
from mylib.mycsv import sort_csv, csv2excel_link


def get_all_stock(idx, stock):
    res_dict = collections.defaultdict(list)
    df = pd.read_csv('all.csv')
    count = 0
    for row in df.index:
        if stock and df.loc[row]['ts_code'] != stock:
            continue
        sn = StockNumber(df.loc[row])
        if '退' in sn.name:
            continue
        if 'ST' in sn.name:
            continue
        if not str(sn.ts_code).startswith('60') \
                and not str(sn.ts_code).startswith('00'):
            continue
        count += 1
        n, msg = analysis_stock(idx, sn)
        if n is None and msg is None:
            continue
        res_dict[n].append(msg)
    return res_dict


def analysis_stock(idx, sn):
    csv_path = f'stocks/{sn.ts_code}.csv'
    if not os.path.exists(csv_path):
        pro = ts.pro_api(token='bbc6f076aa3d2b063cb26376ad12ffd94b53864b42d2344b3aa2039f')
        df2 = pro.daily(ts_code=sn.ts_code)
        save_dir = f'../../stocks'
        if not os.path.exists(save_dir):
            os.mkdir(save_dir)
        df2.to_csv(csv_path)
        print(f'save {sn.ts_code} {sn.name} 成功')
    df2 = pd.read_csv(csv_path)
    down_arr = []
    all_days_data = []
    today_di = None
    N = 100000
    start_arr = []
    for row2 in df2.index:
        if row2 < idx:
            continue
        if N == 1:
            return None, None
        if len(down_arr) == N:
            link_code_arr = sn.ts_code.split('.')
            link_code = f'{link_code_arr[1]}{link_code_arr[0]}'
            hyperlink = f'https://xueqiu.com/S/{link_code},{today_di.name}'
            msg = f'{hyperlink},{today_di.industry},{today_di.trade_date},{today_di.close},{down_arr[0].trade_date},{down_arr[0].close},{len(all_days_data)},天前{N}连涨\n'
            logger.info(msg)
            return N, msg
        di = DayInfo(sn, df2.loc[row2])
        di.name = sn.name
        di.ts_code = sn.ts_code
        di.industry = sn.industry
        # di.trade_date = df2.loc[row2]['trade_date']
        # di.open = df2.loc[row2]['open']
        # di.high = df2.loc[row2]['high']
        # di.low = df2.loc[row2]['low']
        # di.close = float(df2.loc[row2]['close'])
        if di.close < glm.min_price:
            return None, None
        # di.pre_close = df2.loc[row2]['pre_close']
        # di.change = df2.loc[row2]['change']
        # di.pct_chg = float(df2.loc[row2]['pct_chg'])
        if not start_arr and di.pct_chg <= 0:
            return None, None
        # di.vol = df2.loc[row2]['vol']
        # di.amount = df2.loc[row2]['amount']
        if today_di is None:
            today_di = di
        # 计算当前连续下涨次数
        if N == 100000:
            if di.pct_chg > 0:
                start_arr.append(di)
                continue
            else:
                N = len(start_arr)
        if di.pct_chg > 0:
            down_arr.append(di)
        else:
            down_arr = []
        all_days_data.append(di)

    return None, None


def printNumber(idx, stocks):
    res_dict = {}
    for stock in stocks:
        res_dict = get_all_stock(idx, stock)
    return res_dict


def run(stocks, bkd=0):
    logger.info(f'生成csv start')
    for idx, trade_date in enumerate(glm.all_trade_date_arr):
        today_str = str(trade_date)
        if idx > bkd:
            break
        logger.info(f'开始计算{today_str}数据')
        res_dict = printNumber(idx, stocks)
        for n, data_list in res_dict.items():
            csv_path = f'result/{n}_{today_str}_连涨.csv'
            with open(csv_path, 'a+') as fw:
                title = "超链接,名字,行业,今日日期,今日价格,不满足日期,不满足价格,ndays,描述\n"
                fw.write(title)
                fw.writelines(data_list)
            sort_csv(csv_path, 'ndays')
            csv2excel_link(csv_path)
        logger.info(f'{today_str}计算完成')
